# 添加节点
import time

import networkx as nx
import matplotlib.pyplot as plt


def draw(origin_data, data, epoch, idx):
    # print(data)
    # print("batch.node_id...........",origin_data.node_id)
    # # print("batch.node_type...........",origin_data.node_type)
    # # n_id with edge_label_index
    # print("batch.n_id...........", origin_data.n_id)
    # print("batch.edge_label_index...........", origin_data.edge_label_index)
    # print("batch.edge_label...........", origin_data.edge_label)
    # # print(origin_data)
    # a = [origin_data.edge_label_index.t().tolist()]
    # origin_data.node_id.tolist()

    G = nx.Graph()
    nid_start = origin_data.n_id[origin_data.edge_label_index[0]]
    nid_end = origin_data.n_id[origin_data.edge_label_index[1]]

    nid = origin_data.n_id.tolist()
    nodes = list(set(nid_start.tolist() + nid_end.tolist()))

    ###############
    a_id_list = [i for i in range(len(data["alarm"].node_id))]
    h_id_list = [i for i in range(len(data["alarm"].node_id), len(data["host"].node_id) + len(data["alarm"].node_id))]
    b_id_list = [i for i in range(len(data["host"].node_id) + len(data["alarm"].node_id),
                                  len(data["host"].node_id) + len(data["alarm"].node_id) + len(
                                      data["bussiness_tree"].node_id))]

    alarm_node_list = [i for i in nodes if i in a_id_list]
    host_node_list = [i for i in nodes if i in h_id_list]
    btree_node_list = [i for i in nodes if i in b_id_list]

    G.add_nodes_from(alarm_node_list, color='red', size=10)
    G.add_nodes_from(host_node_list, color='green', size=20)
    G.add_nodes_from(btree_node_list, color='blue', size=30)

    edge_true_mask = origin_data.edge_label == 1
    edge_false_mask = origin_data.edge_label == 0

    true_edges = origin_data.edge_label_index.t()[edge_true_mask].tolist()
    false_edges = origin_data.edge_label_index.t()[edge_false_mask].tolist()

    true_edge_list = [[nid[e[0]], nid[e[1]]] for e in true_edges]
    false_edge_list = [[nid[e[0]], nid[e[1]]] for e in false_edges]
    # print("1_nodes", nodes)
    # print("2_true_edge_list", true_edge_list)
    # print("3_false_edge_list", false_edge_list)
    #
    # for i in true_edge_list:
    #     print("4_true_like_data: {}".format(data["all_dict"][i[0]]+"||||"+data["all_dict"][i[1]]))

    G.add_edges_from(true_edge_list, color='red', size=10)
    G.add_edges_from(false_edge_list, color='gray', size=10)

    pos = nx.spring_layout(G, k=1, scale=10)
    # pos =  nx.spring_layout(G,k=2,scale=2)
    # pos = nx.circular_layout(G,scale=1000)
    # pos = nx.shell_layout(G)
    # pos = nx.spectral_layout(G)
    # print("##1111",[n[1]["shape"] if "shape" in n[1] else "8" for n in G.nodes(data=True)])

    nx.draw(G, pos=pos, with_labels=True, node_size=200,
            # style="--",
            # 'so^>v<dph8'
            node_shape="o",
            # font_size=[n[1]["size"] if "size" in n[1] else 10 for n in G.nodes(data=True)],
            font_size=10,
            node_color=[n[1]["color"] if "color" in n[1] else "black" for n in G.nodes(data=True)],
            edge_color=[n[2]["color"] if "color" in n[2] else "black" for n in G.edges(data=True)]
            # width =2.0
            )
    # plt.show()
    # time.sleep(1)
    plt.savefig("/home/yitian001/code/storage_models/logs/output_{}_{}.png".format(epoch, idx))
